# This file is part of Gehyra.
#
# Gehyra is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Gehyra is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Gehyra.  If not, see <http://www.gnu.org/licenses/>.

"""
@package gehyra.common.ipv6
Defines utility classes for dealing with IPv6 address and transforming between
various representations.
$Id: ipv6.py 460 2011-02-11 19:33:54Z andyhhp@gmail.com $
"""

"""
@file gehyra/common/ipv6.py
Defines utility classes for dealing with IPv6 address and transforming between
various representations.
"""

import re, struct


class IPv6Address(object):
    """Class for manipulating IPv6 Addresses.
    Can deal with dotted decimal representation, packed structures
    and integers in string or int form.
    """

    ## regex for finding hexidecimal string representations
    hex_regex = re.compile(r"^(?:0x)?(?P<addr>[0-9a-fA-F]{32})$")

    ## regex for finding the common string representations 
    common_regex_recogniser = re.compile(
        r"\[?[0-9a-fA-F]{0,4}:(?:[0-9a-fA-F]{0,4}:?){0,6}:[0-9a-fA-F]{0,4}\]?$")
    common_regex_extractor = re.compile(r"[0-9a-fA-F]{1,4}")

    def __init__(self, address = None):
        """Constructor.
        """

        ## integer value of the IPv6 address
        self.ip = None

        if isinstance(address, (int, long)):
            self.set_integer(address)

        elif isinstance(address, str):

            if IPv6Address.hex_regex.match(address) is not None:
                self.set_hex_string(address)
                return

            if IPv6Address.common_regex_recogniser.match(address) is not None:
                self.set_common_string(address)
                return

            try:
                self.set_packed_string(address)
            except ValueError:
                pass
            else:
                return
            
            raise ValueError("Unrecognised IPv6 Address '%s'" % (address,))

        else:
            raise TypeError("String or Integer expected.  "
                            "Got %s" % type(address))
        
    def __ne__(self, other):
        """overload the != operator"""

        if not isinstance(other, IPv6Address):
            try:
                other = IPv6Address(other)
            except Exception:
                return NotImplemented

        return self.ip != other.ip

    def __eq__(self, other):
        """overload the == operator"""

        if not isinstance(other, IPv6Address):
            try:
                other = IPv6Address(other)
            except Exception:
                return NotImplemented

        return self.ip == other.ip

    def set_integer(self, addr):
        """Set the IPv6 address using a long value.
        @param addr integer IPv6 address
        @exception TypeError if addr is not a long
        @exception ValueError if addr is outside the range 0 <= val < 2^128
        """

        if not isinstance(addr, (int, long)):
            raise TypeError("Expected Integer.  Got %s" % type(addr))

        if 0 <= addr < 2**128:
            self.ip = long(addr)
        else:
            raise ValueError("IPv6 address must be in range 0 <= val < 2^128")


    def set_common_string(self, addr):
        """Set the IPv6 address using a dotted decimal string.
        @param addr Colon hexidecimal string
        @exception TypeError if addr is not a string
        @exception ValueError if addr is not a valid dotted decimal
        representation or if any of the sextets are outside the range
        0 <= val <= 65535
        """

        if not isinstance(addr, str):
            raise TypeError("Expected String.  Got %s" % type(addr))

        # remove the optional square brackets
        if addr[0] == '[' and addr[-1] == ']':
            addr = addr[1:-1]

        if addr.find('::') != -1:
            # we have a shortened address

            front, back = addr.split('::', 1)

            if (len(back) > 0 and back[0] == ':') or (back.find('::') != -1):
                raise ValueError("'%s' is not a valid contracted IPv6 "
                                 "Address" % addr)
            
            front_quads = IPv6Address.common_regex_extractor.findall(front)
            back_quads = IPv6Address.common_regex_extractor.findall(back)

            implicit_quads = 8 - len(front_quads) - len(back_quads)
            
            if implicit_quads < 1:
                raise ValueError("'%s' is not a valid contracted IPv6 "
                                 "Address" % addr)

            quads = front_quads + (["0"] * implicit_quads) + back_quads


        else:
            # full number of quads expected
            quads = IPv6Address.common_regex_extractor.findall(addr)

            if len(quads) != 8:
                raise ValueError("Expected 8 quad groups")

        octets = []
        for q in quads:
            l = long(q, 16)
            
            octets.extend( [ (l >> 8) & 0xff, l & 0xff ] )
            
        self._set_from_octets(octets)

    def set_hex_string(self, addr):
        """Set the IPv6 address using a hexidecimal string.
        @param addr Hexidecimal string with an optional 0x prefix
        @exception TypeError if addr is not a string
        @exception ValueError if addr is not a valid hexidecimal representation
        or if the hexidecimal number is out outside the range 0 <= val < 2^128
        """

        if not isinstance(addr, str):
            raise TypeError("Expected String.  Got %s" % type(addr))

        res = IPv6Address.hex_regex.match(addr)
        if res is None:
            raise ValueError("Expected Hex value in string.  Got %s" % addr)

        self.set_integer(long(res.group('addr'), 16))

    def set_packed_string(self, addr):
        """Set the IPv6 address using a packed structure.
        @param addr packed structure string
        @exception TypeError if addr is not a string
        @exception ValueError if addr is out of the range 0 <= val < 2^128
        """

        if not isinstance(addr, str):
            raise TypeError("Expected String.  Got %s" % type(addr))

        try:
            self._set_from_octets(struct.unpack("!16B", addr))
        except struct.error:
            raise ValueError("Packed value out of range")
        


    def get_integer(self):
        """Get the IPv6 address as an integer.
        @exception ValueError if self.ip has not been set
        @returns integer IPv6 address
        """

        if self.ip is None:
            raise ValueError("Address has not been set to any value")
        return self.ip

    def get_common_string(self, expanded=False):
        """Get the IPv6 address as a dotted decimal string.
        @param expanded boolean value indicating whether to expand each octet
        string to be 3 characters with leading 0's where nessesary.
        @exception ValueError if self.ip has not been set
        @returns dotted decimal string
        """
        
        if self.ip is None:
            raise ValueError("Address has not been set to any value")
        
        sextets = [ (self.ip >> 112) & 0xffff,
                    (self.ip >> 96)  & 0xffff,
                    (self.ip >> 80)  & 0xffff,
                    (self.ip >> 64)  & 0xffff,
                    (self.ip >> 48)  & 0xffff,
                    (self.ip >> 32)  & 0xffff,
                    (self.ip >> 16)  & 0xffff,
                    (self.ip)        & 0xffff ]
        
        if expanded:
            return ':'.join( "%04X" % s for s in sextets )
        else:
            return ':'.join( "%X" % s for s in sextets )

    def get_hex_string(self, prefix=False):
        """Get the IPv6 address as an hexidecmal string.
        @param prefix boolean value indicating whether a 0x prefix is wanted on
        the returned string
        @exception ValueError if self.ip has not been set
        @returns hexidecimal string
        """
        
        if self.ip is None:
            raise ValueError("Address has not been set to any value")
        
        if prefix:
            return "0x%032X" % self.ip
        else:
            return "%032X" % self.ip


    def get_packed_string(self):
        """Get the IPv6 address as a packed string in network order.
        @exception ValueError if self.ip has not been set
        @returns packed string
        """

        if self.ip is None:
            raise ValueError("Address has not been set to any value")

        return struct.pack("!16B", *self._get_as_octets())


    def _set_from_octets(self, octets):
        """Private helper to set self.ip from a list of octets"""
        self.ip = long( (octets[ 0] << 120) |
                        (octets[ 1] << 112) |
                        (octets[ 2] << 104) |
                        (octets[ 3] <<  96) |
                        (octets[ 4] <<  88) |
                        (octets[ 5] <<  80) |
                        (octets[ 6] <<  72) |
                        (octets[ 7] <<  64) |
                        (octets[ 8] <<  56) |
                        (octets[ 9] <<  48) |
                        (octets[10] <<  40) |
                        (octets[11] <<  32) |
                        (octets[12] <<  24) |
                        (octets[13] <<  16) |
                        (octets[14] <<   8) |
                        (octets[15]) )

    def _get_as_octets(self):
        """Private helper to get a list of octets from self.ip"""
        return [ (self.ip >> 120) & 0xff,
                 (self.ip >> 112) & 0xff,
                 (self.ip >> 104) & 0xff,
                 (self.ip >>  96) & 0xff,
                 (self.ip >>  88) & 0xff,
                 (self.ip >>  80) & 0xff,
                 (self.ip >>  72) & 0xff,
                 (self.ip >>  64) & 0xff,
                 (self.ip >>  56) & 0xff,
                 (self.ip >>  48) & 0xff,
                 (self.ip >>  40) & 0xff,
                 (self.ip >>  32) & 0xff,
                 (self.ip >>  24) & 0xff,
                 (self.ip >>  16) & 0xff,
                 (self.ip >>   8) & 0xff,
                 (self.ip       ) & 0xff]
